#include <iostream>
#include <Eigen/Sparse>
#include "CFEMNonlinear.h"
#include <time.h>
#define PI 3.1415926535897


double nu(CPoint p, double t)
{
	return 1;
}
double  f1(CPoint p, double t)
{
	return -2 * nu(p, t)*pow(p.getX(), 2) - 2 * nu(p, t)*pow(p.getY(), 2) - nu(p, t)*exp(-p.getY()) + pow(PI, 2)*cos(PI*p.getX())*cos(2 * PI*p.getY())
		+2*p.getX()*pow(p.getY(),2)*(pow(p.getX(),2)*pow(p.getY(),2)+exp(-p.getY()))+(-2*p.getX()*pow(p.getY(),3)/3.0+2-PI*sin(PI*p.getX()))*(2*pow(p.getX(),2)*p.getY()-exp(-p.getY()));
}
double f2(CPoint p, double t)
{
	return 4 * nu(p, t)*p.getX()*p.getY() - nu(p, t)*pow(PI, 3)*sin(PI*p.getX()) + 2 * PI*(2 - PI * sin(PI*p.getX()))*sin(2 * PI*p.getY())
		+(pow(p.getX(),2)*pow(p.getY(),2)+exp(-p.getY()))*(-2*pow(p.getY(),3)/3.0-pow(PI,2)*cos(PI*p.getX()))
		+(-2*p.getX()*pow(p.getY(),3)/3.0+2-PI*sin(PI*p.getX()))*(-2*p.getX()*pow(p.getY(),2));
}
double u1_left(CPoint p, double t)
{
	return exp(-p.getY());
}
double u1_right(CPoint p, double t)
{
	return pow(p.getY(), 2) + exp(-p.getY());
}
double u1_bottom(CPoint p, double t)
{
	return 1.0 / 16.0*pow(p.getX(), 2) + exp(0.25);
}
double u1_up(CPoint p, double t)
{
	return 1;
}
double u2_left(CPoint p, double t)
{
	return 2;
}
double u2_right(CPoint p, double t)
{
	return -2.0 / 3.0*pow(p.getY(), 3) + 2;
}
double u2_bottom(CPoint p, double t)
{
	return 1.0 / 96.0*p.getX() + 2 - PI * sin(PI*p.getX());
}
double u2_up(CPoint p, double t)
{
	return 2 - PI * sin(PI*p.getX());
}
double fun_negative_one(CPoint, double)
{
	return -1.0;
}

double coe_fun(CPoint p, double t)
{
	return 1.0;
}
int main()
{
	int unknown_size = 3;
	CFEMNonlinear fem(unknown_size);
	//read mesh from file
	//the mesh file contains the P,T,Pb,Tb,boundary nodes,boundary edges, and the type of the element
	fem.readMesh("..\\mesh\\cha6_exa1_4.neu");

	//The initializing process need mesh's information, so the function readMesh should be called first
	fem.Init();

	fem.assignBoundNodesValue(0, 0, Dirichlet, u1_left);
	fem.assignBoundNodesValue(1, 0, Dirichlet, u1_bottom);
	fem.assignBoundNodesValue(2, 0, Dirichlet, u1_right);
	fem.assignBoundNodesValue(3, 0, Dirichlet, u1_up);

	fem.assignBoundNodesValue(0, 1, Dirichlet, u2_left);
	fem.assignBoundNodesValue(1, 1, Dirichlet, u2_bottom);
	fem.assignBoundNodesValue(2, 1, Dirichlet, u2_right);
	fem.assignBoundNodesValue(3, 1, Dirichlet, u2_up);

	CBasisFunction u_basis_trial(TWO_DIM_QUADRATIC);
	CBasisFunction u_basis_test(TWO_DIM_QUADRATIC);

	CBasisFunction p_basis_trial(TWO_DIM_LINEAR);
	CBasisFunction p_basis_test(TWO_DIM_LINEAR);

	pair<int, int> matrix_size_u{ fem.getPbTest().size(),fem.getPbTrial().size() };
	pair<int, int> matrix_size_p{ fem.getPbTest().size(),fem.getP().size() };

	//Maximum number of Newton iteration steps
	int max_iteration = 10;

	//convergence tolerance
	//it means that if the difference between the current iterated step and previous one is less than the tol,
	//then the iteration is convergent even the iteration step doesn't reach the maximum.
	double tol = 1.0e-5;

	//initial guess of the unknowns.
	//we let them to be zero, it is should be noticed that this is just for the demo example.
	//in the real physical problem, it is better to set the velocity and pressure to be reasonable initial values.
	Rsv u1_0(matrix_size_u.first), u1_1(matrix_size_u.first);
	Rsv u2_0(matrix_size_u.first), u2_1(matrix_size_u.first);
	//Rsv p_0(matrix_size_p.second), p_1(matrix_size_p.second);
	u1_0.setZero();
	u2_0.setZero();
	//p_0.setZero();
	
	//iterating step
	int l = 1;

	while (true)	
	{		
		cout << "\nIterating step:\t" << l<<endl;
		//assemble matrix of the linear terms of the NS equations.
		SpMat A1 = fem.assembleMatrix2D(nu, 0, u_basis_trial, 1, 0, fem.getTbTrial(), u_basis_test, 1, 0, fem.getTbTest(), matrix_size_u);
		SpMat A2 = fem.assembleMatrix2D(nu, 0, u_basis_trial, 0, 1, fem.getTbTrial(), u_basis_test, 0, 1, fem.getTbTest(), matrix_size_u);
		SpMat A3 = fem.assembleMatrix2D(nu, 0, u_basis_trial, 1, 0, fem.getTbTrial(), u_basis_test, 0, 1, fem.getTbTest(), matrix_size_u);
		SpMat A5 = fem.assembleMatrix2D(fun_negative_one, 0, p_basis_trial, 0, 0, fem.getT(), u_basis_test, 1, 0, fem.getTbTest(), matrix_size_p);
		SpMat A6 = fem.assembleMatrix2D(fun_negative_one, 0, p_basis_trial, 0, 0, fem.getT(), u_basis_test, 0, 1, fem.getTbTest(), matrix_size_p);
		SpMat zero(matrix_size_p.second, matrix_size_p.second);
		zero.setZero();
		SpMat A = fem.matrix_vstack(fem.matrix_hstack(2 * A1 + A2, A3, A5), fem.matrix_hstack(A3.transpose(), 2 * A2 + A1, A6), fem.matrix_hstack(A5.transpose(), A6.transpose(), zero));

		//assemble vector of the linear term of the NS equations.
		Rsv b1 = fem.assembleVector2D(f1, 0, u_basis_test, 0, 0, fem.getTbTest());
		Rsv b2 = fem.assembleVector2D(f2, 0, u_basis_test, 0, 0, fem.getTbTest());
		Rsv zero_vec(fem.getP().size());
		zero_vec.setZero();
		Rsv b = fem.vector_stack(b1, b2, zero_vec);

		//assemble matrix for the nonlinear terms of the NS equations.
		SpMat AN1 = fem.assembleMatrix2DNonlinear(
			u1_0, 0, u_basis_trial, 1, 0, fem.getTbTrial(), //coefficient
			u_basis_trial, 0, 0, fem.getTbTrial(), u_basis_test, 0, 0, fem.getTbTest(), matrix_size_u);
		SpMat AN2 = fem.assembleMatrix2DNonlinear(
			u1_0, 0, u_basis_trial, 0, 0, fem.getTbTrial(), //coefficient
			u_basis_trial, 1, 0, fem.getTbTrial(), u_basis_test, 0, 0, fem.getTbTest(), matrix_size_u);
		SpMat AN3 = fem.assembleMatrix2DNonlinear(
			u2_0, 0, u_basis_trial, 0, 0, fem.getTbTrial(), //coefficient
			u_basis_trial, 0, 1, fem.getTbTrial(), u_basis_test, 0, 0, fem.getTbTest(), matrix_size_u);
		SpMat AN4 = fem.assembleMatrix2DNonlinear(
			u1_0, 0, u_basis_trial, 0, 1, fem.getTbTrial(), //coefficient
			u_basis_trial, 0, 0, fem.getTbTrial(), u_basis_test, 0, 0, fem.getTbTest(), matrix_size_u);
		SpMat AN5 = fem.assembleMatrix2DNonlinear(
			u2_0, 0, u_basis_trial, 1, 0, fem.getTbTrial(), //coefficient
			u_basis_trial, 0, 0, fem.getTbTrial(), u_basis_test, 0, 0, fem.getTbTest(), matrix_size_u);
		SpMat AN6 = fem.assembleMatrix2DNonlinear(
			u2_0, 0, u_basis_trial, 0, 1, fem.getTbTrial(), //coefficient
			u_basis_trial, 0, 0, fem.getTbTrial(), u_basis_test, 0, 0, fem.getTbTest(), matrix_size_u);
		SpMat zero_1(matrix_size_p.first, matrix_size_p.second);
		zero_1.setZero();
		SpMat AN = fem.matrix_vstack(fem.matrix_hstack(AN1 + AN2 + AN3, AN4, zero_1), fem.matrix_hstack(AN5, AN6 + AN2 + AN3, zero_1), fem.matrix_hstack(zero_1.transpose(), zero_1.transpose(), zero));
			   		 		
		//assemble vector for the nonlinear terms of NS equations.
		Rsv bN1 = fem.assembleVector2DNonlinear(0,
			u1_0, u_basis_trial, 0, 0, fem.getTbTrial(), //coefficient1
			u1_0, u_basis_trial, 1, 0, fem.getTbTrial(), //coefficient2
			u_basis_test, 0, 0, fem.getTbTest());
		Rsv bN2 = fem.assembleVector2DNonlinear(0,
			u2_0, u_basis_trial, 0, 0, fem.getTbTrial(),//coefficient1
			u1_0, u_basis_trial, 0, 1, fem.getTbTrial(),//coefficient2
			u_basis_test, 0, 0, fem.getTbTest());
		Rsv bN3 = fem.assembleVector2DNonlinear(0,
			u1_0, u_basis_trial, 0, 0, fem.getTbTrial(),//coefficient1
			u2_0, u_basis_trial, 1, 0, fem.getTbTrial(),//coefficient2
			u_basis_test, 0, 0, fem.getTbTest());
		Rsv bN4 = fem.assembleVector2DNonlinear(0,
			u2_0, u_basis_trial, 0, 0, fem.getTbTrial(),//coefficient1
			u2_0, u_basis_trial, 0, 1, fem.getTbTrial(),//coefficient2
			u_basis_test, 0, 0, fem.getTbTest());
		Rsv bN = fem.vector_stack(bN1 + bN2, bN3 + bN4, zero_vec);

		//the total matrix and vector of the current iteration step
		SpMat Al = A + AN;
		Rsv bl = b + bN;

		//treat boundary conditions
		fem.treatBoundaryConditions(coe_fun, 0, Al, bl, vector<unsigned int>{fem.getPbTrial().size(), fem.getPbTrial().size(), fem.getP().size()});
		//fem.ExportMat("A.txt", A);

		fem.imposeFixValue(1 + 2 *matrix_size_u.first, 0, Al, bl);

		Rsv x = fem.solveLinearEqua(DIREC_PardisoLU, Al, bl);

		u1_1 = x.segment(0, matrix_size_u.first);
		u2_1 = x.segment(matrix_size_u.first, matrix_size_u.first);
		
		double error1 = fabs(u1_1(0) - u1_0(0));
		for (size_t i = 0; i < u1_1.size(); i++)
		{
			if (error1 < fabs(u1_1(i) - u1_0(i)))
			{
				error1 = fabs(u1_1(i) - u1_0(i));
			}
		}
		double error2 = fabs(u2_1(0) - u2_0(0));
		for (size_t i = 0; i < u2_1.size(); i++)
		{
			if (error2 < fabs(u2_1(i) - u2_0(i)))
			{
				error2 = fabs(u2_1(i) - u2_0(i));
			}
		}
		cout << "max error of the u1:\t" << error1 << endl;
		cout << "max error of the u2:\t" << error2 << endl;

		if (error1 <= tol && error2 <= tol)
		{
			cout << "\nStop:Convergent" << endl;
			fem.export2VTK("u.vtk", x);

			Rsv x_p(fem.getP().size());
			for (int i = 0; i < x_p.rows(); ++i)
			{
				x_p(i) = x(i + 2 * fem.getPbTest().size());
			}

			fem.export2VTK("p.vtk", x_p, fem.getP(), fem.getT());
			break;//stop the iteration
		}
		else//continue the iteration
		{
			++l;
			if (l==max_iteration)
			{
				cout << "\nStop: Reach the maximum iteration steps and not convergent.Try another initial guess or increase the maximum iteration step" << endl;
				break;
			}		
			cout << "\nNot convergent, begin next iteration" << endl;

			u1_0 = u1_1;
			u2_0 = u2_1;
		}
		
	}

	
}
